#!/usr/bin/env python
# @Date    : 2023-03-29
# @Author  : Bright (brt2@qq.com)
# @Link    : https://gitee.com/brt2
# @Version : 0.1.0

import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

from .cvs import pimg2array

def show_image(data_item):
    img, target = data_item
    # print('Label: ', datasets.classes[target])
    plt.imshow(img.permute(1,2,0))

def imshow(im, setHD=False):
    if isinstance(im, Image.Image):
        npimg = pimg2array(im)
    elif isinstance(im, np.ndarray):  # torch.Tensor
        npimg = im
    else:
        # npimg = np.transpose(im.numpy(), (1, 2, 0))
        npimg = im.permute((1, 2, 0))

    if setHD:
        plt.figure(figsize=(20,20))
    plt.imshow(npimg)

def cvtColor(im, *args, **kwargs):
    pass

# %% 显示图像
def show_images_batch(dataloader):
    import torchvision

    for images, _labels in dataloader:
        _fig, ax= plt.subplots(figsize=(16,8))
        ax.set_xticks([])
        ax.set_yticks([])
        ax.imshow(torchvision.utils.make_grid(images, nrow=16).permute(1,2,0))
        break  # to stop loop otherwise 4500 images in batch size of 128 will print and is computationally expensive
